1use crate::ext::io::*;
3use crate::scripts::base::*;
4use crate::types::*;
5use crate::utils::encoding::*;
6use crate::utils::escape::*;
7use anyhow::Result;
8use serde::{Deserialize, Serialize};
9use std::collections::BTreeMap;
10use std::io::{Read, Write};
11use std::ops::Index;
12use unicode_segmentation::UnicodeSegmentation;
13
14#[derive(Debug)]
15pub struct ArtemisAsbBuilder {}
17
18impl ArtemisAsbBuilder {
19 pub fn new() -> Self {
21 ArtemisAsbBuilder {}
22 }
23}
24
25impl ScriptBuilder for ArtemisAsbBuilder {
26 fn default_encoding(&self) -> Encoding {
27 Encoding::Utf8
28 }
29
30 fn build_script(
31 &self,
32 buf: Vec<u8>,
33 _filename: &str,
34 encoding: Encoding,
35 _archive_encoding: Encoding,
36 config: &ExtraConfig,
37 _archive: Option<&Box<dyn Script>>,
38 ) -> Result<Box<dyn Script>> {
39 Ok(Box::new(Asb::new(buf, encoding, config)?))
40 }
41
42 fn extensions(&self) -> &'static [&'static str] {
43 &["asb"]
44 }
45
46 fn script_type(&self) -> &'static ScriptType {
47 &ScriptType::ArtemisAsb
48 }
49
50 fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
51 if buf_len >= 5 && buf.starts_with(b"ASB\0\0") {
52 return Some(20);
53 }
54 None
55 }
56
57 fn can_create_file(&self) -> bool {
58 true
59 }
60
61 fn create_file<'a>(
62 &'a self,
63 filename: &'a str,
64 writer: Box<dyn WriteSeek + 'a>,
65 encoding: Encoding,
66 file_encoding: Encoding,
67 config: &ExtraConfig,
68 ) -> Result<()> {
69 create_file(
70 filename,
71 writer,
72 encoding,
73 file_encoding,
74 config.custom_yaml,
75 )
76 }
77}
78
79fn escape_text(s: &str) -> String {
80 let mut escaped = String::with_capacity(s.len());
81 for c in s.chars() {
82 match c {
83 '&' => escaped.push_str("&"),
84 '<' => escaped.push_str("<"),
85 _ => escaped.push(c),
86 }
87 }
88 escaped
89}
90
91#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
92struct Command {
93 pub name: String,
94 pub line_number: u32,
95 pub attributes: BTreeMap<String, String>,
96}
97
98impl Command {
99 pub fn new(name: String, line_number: u32) -> Self {
100 Command {
101 name,
102 line_number,
103 attributes: BTreeMap::new(),
104 }
105 }
106
107 pub fn to_xml(&self) -> String {
108 let mut xml = format!("<{}", self.name);
109 for (key, value) in &self.attributes {
110 xml.push_str(&format!(" {}=\"{}\"", key, escape_xml_text_value(value)));
111 }
112 xml.push('>');
113 xml
114 }
115}
116
117impl<'a> Index<&'a str> for Command {
118 type Output = str;
119 fn index(&self, key: &'a str) -> &Self::Output {
120 self.attributes.get(key).map_or("", |s| s.as_str())
121 }
122}
123
124impl<'a> Index<&'a String> for Command {
125 type Output = str;
126 fn index(&self, key: &'a String) -> &Self::Output {
127 self.attributes.get(key).map_or("", |s| s.as_str())
128 }
129}
130
131impl Index<String> for Command {
132 type Output = str;
133 fn index(&self, key: String) -> &Self::Output {
134 self.attributes.get(&key).map_or("", |s| s.as_str())
135 }
136}
137
138#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
139#[serde(untagged)]
140enum Item {
141 Command(Command),
142 Label(String),
143}
144
145impl Item {
146 pub fn is_command(&self) -> bool {
147 matches!(self, Item::Command(_))
148 }
149
150 pub fn is_command_name(&self, name: &str) -> bool {
151 if let Item::Command(cmd) = self {
152 cmd.name == name
153 } else {
154 false
155 }
156 }
157}
158
159trait CustomReadFn {
160 fn read_string(&mut self, encoding: Encoding) -> Result<String>;
161 fn read_item(&mut self, encoding: Encoding) -> Result<Item>;
162}
163
164impl<T: Read> CustomReadFn for T {
165 fn read_string(&mut self, encoding: Encoding) -> Result<String> {
166 let len = self.read_u32()?;
167 let data = self.read_exact_vec(len as usize)?;
168 if self.read_u8()? != 0 {
169 return Err(anyhow::anyhow!("String not null-terminated"));
170 }
171 let s = decode_to_string(encoding, &data, true)?;
172 Ok(s)
173 }
174
175 fn read_item(&mut self, encoding: Encoding) -> Result<Item> {
176 let typ = self.read_u32()?;
177 match typ {
178 0 => {
179 let name = self.read_string(encoding)?;
180 let line_number = self.read_u32()?;
181 let mut command = Command::new(name, line_number);
182 let attr_count = self.read_u32()?;
183 for _ in 0..attr_count {
184 let key = self.read_string(encoding)?;
185 let value = self.read_string(encoding)?;
186 command.attributes.insert(key, value);
187 }
188 Ok(Item::Command(command))
189 }
190 1 => {
191 let label = self.read_string(encoding)?;
192 Ok(Item::Label(label))
193 }
194 _ => {
195 return Err(anyhow::anyhow!("Unknown item type: {}", typ));
196 }
197 }
198 }
199}
200
201trait CustomWriteFn {
202 fn write_string(&mut self, s: &str, encoding: Encoding) -> Result<()>;
203 fn write_item(&mut self, item: &Item, encoding: Encoding) -> Result<()>;
204}
205
206impl<T: Write> CustomWriteFn for T {
207 fn write_string(&mut self, s: &str, encoding: Encoding) -> Result<()> {
208 let data = encode_string(encoding, s, false)?;
209 self.write_u32(data.len() as u32)?;
210 self.write_all(&data)?;
211 self.write_u8(0)?; Ok(())
213 }
214
215 fn write_item(&mut self, item: &Item, encoding: Encoding) -> Result<()> {
216 match item {
217 Item::Command(cmd) => {
218 self.write_u32(0)?; self.write_string(&cmd.name, encoding)?;
220 self.write_u32(cmd.line_number)?;
221 self.write_u32(cmd.attributes.len() as u32)?;
222 for (key, value) in &cmd.attributes {
223 self.write_string(key, encoding)?;
224 self.write_string(value, encoding)?;
225 }
226 }
227 Item::Label(label) => {
228 self.write_u32(1)?; self.write_string(label, encoding)?;
230 }
231 }
232 Ok(())
233 }
234}
235
236struct TextParser<'a> {
237 items: Vec<Item>,
238 text: Vec<&'a str>,
239 pos: usize,
240 len: usize,
241 hcls_index: usize,
242}
243
244impl<'a> TextParser<'a> {
245 pub fn new(str: &'a str, hcls_index: usize) -> Self {
246 let text: Vec<&'a str> = UnicodeSegmentation::graphemes(str, true).collect();
247 let len = text.len();
248 TextParser {
249 items: Vec::new(),
250 text,
251 pos: 0,
252 len,
253 hcls_index,
254 }
255 }
256
257 pub fn parse(mut self) -> Result<Vec<Item>> {
258 while let Some(c) = self.peek() {
259 match c {
260 "<" => {
261 self.parse_tag()?;
262 }
263 _ => {
264 let mut text = String::new();
265 self.eat_char();
266 text.push_str(c);
267 while let Some(b) = self.peek() {
268 if b == "<" {
269 break;
270 }
271 text.push_str(b);
272 self.eat_char();
273 }
274 if !text.is_empty() {
275 self.items.push(Item::Command(Command {
276 name: "print".to_string(),
277 line_number: 0,
278 attributes: [("data".to_string(), unescape_xml(&text))].into(),
279 }))
280 }
281 }
282 }
283 }
284 let mut hcls = Command::new("hcls".to_string(), 0);
285 hcls.attributes
286 .insert("0".to_string(), self.hcls_index.to_string());
287 self.items.push(Item::Command(hcls));
288 Ok(self.items)
289 }
290
291 fn parse_tag(&mut self) -> Result<()> {
292 self.parse_indent("<")?;
293 let key = self.parse_key()?;
294 self.erase_whitespace();
295 let mut cmd = Command::new(key, 0);
296 loop {
297 let c = self.peek().ok_or(self.error2("Unexpected eof"))?;
298 match c {
299 ">" => {
300 self.eat_char();
301 break;
302 }
303 " " => {
304 self.eat_char();
305 continue;
306 }
307 _ => {
308 let key = self.parse_key()?;
309 self.parse_indent("=")?;
310 let value = self.parse_str()?;
311 cmd.attributes.insert(key, value);
312 }
313 }
314 }
315 self.items.push(Item::Command(cmd));
316 Ok(())
317 }
318
319 fn parse_key(&mut self) -> Result<String> {
320 self.erase_whitespace();
321 let mut key = String::new();
322 while let Some(c) = self.peek() {
323 if c == "=" || c == " " || c == ">" {
324 break;
325 }
326 key.push_str(c);
327 self.eat_char();
328 }
329 if key.is_empty() {
330 return self.error("Expected key, but found nothing");
331 }
332 Ok(key)
333 }
334
335 fn parse_str(&mut self) -> Result<String> {
336 self.erase_whitespace();
337 self.parse_indent("\"")?;
338 let mut text = String::new();
339 loop {
340 match self.next().ok_or(self.error2("Unexpected eof"))? {
341 "\"" => {
342 break;
343 }
344 t => {
345 text.push_str(t);
346 }
347 }
348 }
349 Ok(unescape_xml(&text))
350 }
351
352 fn erase_whitespace(&mut self) {
353 while let Some(c) = self.peek() {
354 if c == " " {
355 self.eat_char();
356 } else {
357 break;
358 }
359 }
360 }
361
362 fn parse_indent(&mut self, indent: &str) -> Result<()> {
363 for ident in indent.graphemes(true) {
364 match self.next() {
365 Some(c) => {
366 if c != ident {
367 return self.error("Unexpected indent");
368 }
369 }
370 None => return self.error("Unexpected eof"),
371 }
372 }
373 Ok(())
374 }
375
376 fn eat_char(&mut self) {
377 if self.pos < self.len {
378 self.pos += 1;
379 }
380 }
381
382 fn next(&mut self) -> Option<&'a str> {
383 if self.pos < self.len {
384 let item = self.text[self.pos];
385 self.pos += 1;
386 Some(item)
387 } else {
388 None
389 }
390 }
391
392 fn peek(&self) -> Option<&'a str> {
393 if self.pos < self.len {
394 Some(self.text[self.pos])
395 } else {
396 None
397 }
398 }
399
400 fn error2<T>(&self, msg: T) -> anyhow::Error
401 where
402 T: std::fmt::Display,
403 {
404 anyhow::anyhow!("Failed to parse at position {}: {}", self.pos, msg)
405 }
406
407 fn error<T, A>(&self, msg: T) -> Result<A>
408 where
409 T: std::fmt::Display,
410 {
411 Err(anyhow::anyhow!(
412 "Failed to parse at position {}: {}",
413 self.pos,
414 msg
415 ))
416 }
417}
418
419#[derive(Debug)]
420pub struct Asb {
422 items: Vec<Item>,
423 custom_yaml: bool,
424}
425
426impl Asb {
427 pub fn new(buf: Vec<u8>, encoding: Encoding, config: &ExtraConfig) -> Result<Self> {
433 let mut data = MemReader::new(buf);
434 let mut magic = [0; 5];
435 data.read_exact(&mut magic)?;
436 if &magic != b"ASB\0\0" {
437 return Err(anyhow::anyhow!("Invalid ASB magic number: {:?}", magic));
438 }
439 let nums = data.read_u32()?;
440 let mut items = Vec::with_capacity(nums as usize);
441 for _ in 0..nums {
442 items.push(data.read_item(encoding)?);
443 }
444 Ok(Asb {
445 items,
446 custom_yaml: config.custom_yaml,
447 })
448 }
449}
450
451impl Script for Asb {
452 fn default_output_script_type(&self) -> OutputScriptType {
453 OutputScriptType::Json
454 }
455
456 fn is_output_supported(&self, _: OutputScriptType) -> bool {
457 true
458 }
459
460 fn default_format_type(&self) -> FormatOptions {
461 FormatOptions::None
462 }
463
464 fn custom_output_extension<'a>(&'a self) -> &'a str {
465 if self.custom_yaml { "yaml" } else { "json" }
466 }
467
468 fn extract_messages(&self) -> Result<Vec<Message>> {
469 let mut messages = Vec::new();
470 let mut name = None;
471 let mut cur_mes = String::new();
472 let mut in_print = false;
473 for item in self.items.iter() {
474 if in_print {
475 if let Item::Command(cmd) = item {
476 match cmd.name.as_str() {
477 "hcls" => {
478 in_print = false;
479 messages.push(Message {
480 name: name.take(),
481 message: cur_mes,
482 });
483 cur_mes = String::new();
484 }
485 "print" => {
486 cur_mes.push_str(&escape_text(&cmd["data"]));
487 }
488 "rt" => {
489 cur_mes.push('\n');
490 }
491 _ => {
492 cur_mes.push_str(&cmd.to_xml());
493 }
494 }
495 continue;
496 }
497 }
498 if let Item::Command(cmd) = item {
499 match cmd.name.as_str() {
500 "print" => {
501 cur_mes.push_str(&escape_text(&cmd["data"]));
502 in_print = true;
503 }
504 "name" => {
505 let v = (cmd.attributes.len() - 1).to_string();
506 name = Some(cmd[v].to_owned());
507 }
508 "sel_text" => {
509 let t = &cmd["text"];
510 if !t.is_empty() {
511 messages.push(Message {
512 name: None,
513 message: t.to_owned(),
514 });
515 }
516 }
517 "RegisterTextToHistory" => {
518 let t = &cmd["1"];
519 if !t.is_empty() {
520 messages.push(Message {
521 name: None,
522 message: t.to_owned(),
523 });
524 }
525 }
526 _ => {}
527 }
528 }
529 }
530 if !cur_mes.is_empty() {
531 messages.push(Message {
532 name: name.take(),
533 message: cur_mes,
534 });
535 }
536 Ok(messages)
537 }
538
539 fn import_messages<'a>(
540 &'a self,
541 messages: Vec<Message>,
542 mut file: Box<dyn WriteSeek + 'a>,
543 _filename: &str,
544 encoding: Encoding,
545 replacement: Option<&'a ReplacementTable>,
546 ) -> Result<()> {
547 file.write_all(b"ASB\0\0")?;
548 let mut items = self.items.clone();
549 let mut name_index = None;
550 let mut mes_index = 0;
551 let mut item_index = 0;
552 let mut print_index = None;
553 let mut hcls_index = 1;
554 while item_index < items.len() {
555 if let Some(print_ind) = print_index.clone() {
556 if items[item_index].is_command_name("hcls") {
557 let message = messages
558 .get(mes_index)
559 .ok_or(anyhow::anyhow!("Not enough messages."))?;
560 if let Some(name_index) = name_index.take() {
561 let mut name = match &message.name {
562 Some(name) => name.to_owned(),
563 None => return Err(anyhow::anyhow!("Message without name.")),
564 };
565 if let Some(replacement) = replacement {
566 for (k, v) in &replacement.map {
567 name = name.replace(k, v);
568 }
569 }
570 if let Item::Command(cmd) = &mut items[name_index] {
571 if cmd.attributes.len() > 1 {
572 cmd.attributes
573 .insert(format!("{}", cmd.attributes.len() - 1), name);
574 } else {
575 let oname = cmd
576 .attributes
577 .get("0")
578 .ok_or(anyhow::anyhow!("No name attribute found."))?;
579 if oname != &name {
580 cmd.attributes.insert("1".to_string(), name);
581 }
582 }
583 }
584 }
585 let mut m = message.message.clone();
586 if let Some(replacement) = replacement {
587 for (k, v) in &replacement.map {
588 m = m.replace(k, v);
589 }
590 }
591 let new_cmds = TextParser::new(&m.replace("\n", "<rt>"), hcls_index).parse()?;
592 hcls_index += 1;
593 let new_cmds_len = new_cmds.len();
594 items.splice(print_ind..=item_index, new_cmds);
595 print_index = None;
596 item_index = print_ind + new_cmds_len;
597 mes_index += 1;
598 continue;
599 } else if items[item_index].is_command() {
600 item_index += 1;
601 continue;
602 }
603 }
604 if let Item::Command(cmd) = &mut items[item_index] {
605 match cmd.name.as_str() {
606 "print" => {
607 print_index = Some(item_index);
608 }
609 "name" => {
610 name_index = Some(item_index);
611 }
612 "sel_text" => {
613 let message = messages
614 .get(mes_index)
615 .ok_or(anyhow::anyhow!("Not enough messages."))?;
616 let mut m = message.message.clone();
617 if let Some(replacement) = replacement {
618 for (k, v) in &replacement.map {
619 m = m.replace(k, v);
620 }
621 }
622 cmd.attributes.insert("text".to_string(), m);
623 mes_index += 1;
624 }
625 "RegisterTextToHistory" => {
626 let message = messages
627 .get(mes_index)
628 .ok_or(anyhow::anyhow!("Not enough messages."))?;
629 let mut m = message.message.clone();
630 if let Some(replacement) = replacement {
631 for (k, v) in &replacement.map {
632 m = m.replace(k, v);
633 }
634 }
635 cmd.attributes.insert("1".to_string(), m);
636 mes_index += 1;
637 }
638 _ => {}
639 }
640 }
641 item_index += 1;
642 }
643 if mes_index != messages.len() {
644 return Err(anyhow::anyhow!(
645 "Not all messages were processed, expected {}, got {}",
646 messages.len(),
647 mes_index
648 ));
649 }
650 file.write_u32(items.len() as u32)?;
651 for item in items {
652 file.write_item(&item, encoding)?;
653 }
654 file.flush()?;
655 Ok(())
656 }
657
658 fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
659 let s = if self.custom_yaml {
660 serde_yaml_ng::to_string(&self.items)?
661 } else {
662 serde_json::to_string_pretty(&self.items)?
663 };
664 let s = encode_string(encoding, &s, false)?;
665 let mut file = std::fs::File::create(filename)?;
666 file.write_all(&s)?;
667 Ok(())
668 }
669
670 fn custom_import<'a>(
671 &'a self,
672 custom_filename: &'a str,
673 file: Box<dyn WriteSeek + 'a>,
674 encoding: Encoding,
675 output_encoding: Encoding,
676 ) -> Result<()> {
677 create_file(
678 custom_filename,
679 file,
680 encoding,
681 output_encoding,
682 self.custom_yaml,
683 )
684 }
685}
686
687pub fn create_file<'a>(
695 custom_filename: &'a str,
696 mut writer: Box<dyn WriteSeek + 'a>,
697 encoding: Encoding,
698 output_encoding: Encoding,
699 yaml: bool,
700) -> Result<()> {
701 let f = crate::utils::files::read_file(custom_filename)?;
702 let s = decode_to_string(output_encoding, &f, true)?;
703 let items: Vec<Item> = if yaml {
704 serde_yaml_ng::from_str(&s)?
705 } else {
706 serde_json::from_str(&s)?
707 };
708 writer.write_all(b"ASB\0\0")?;
709 writer.write_u32(items.len() as u32)?;
710 for item in items {
711 writer.write_item(&item, encoding)?;
712 }
713 Ok(())
714}
715
716#[test]
717fn test_parse() {
718 let text = "Hello < & World!<tag><tags x=\"123\"><name 0=\"Ok\">Test";
719 let parser = TextParser::new(text, 1);
720 let items = parser.parse().unwrap();
721 assert_eq!(
722 items,
723 vec![
724 Item::Command(Command {
725 name: "print".to_string(),
726 line_number: 0,
727 attributes: [("data".to_string(), "Hello < & World!".to_string())].into(),
728 }),
729 Item::Command(Command {
730 name: "tag".to_string(),
731 line_number: 0,
732 attributes: BTreeMap::new(),
733 }),
734 Item::Command(Command {
735 name: "tags".to_string(),
736 line_number: 0,
737 attributes: [("x".to_string(), "123".to_string())].into(),
738 }),
739 Item::Command(Command {
740 name: "name".to_string(),
741 line_number: 0,
742 attributes: [("0".to_string(), "Ok".to_string())].into(),
743 }),
744 Item::Command(Command {
745 name: "print".to_string(),
746 line_number: 0,
747 attributes: [("data".to_string(), "Test".to_string())].into(),
748 }),
749 Item::Command(Command {
750 name: "hcls".to_string(),
751 line_number: 0,
752 attributes: BTreeMap::from([("0".to_string(), "1".to_string())]),
753 }),
754 ]
755 )
756}